use super::images::*;
use crate::error;

pub trait Pe {
    fn image(&self) -> &[u8];

    fn dos_header(&self) -> &IMAGE_DOS_HEADER {
        unsafe { &*(self.image().as_ptr() as *const IMAGE_DOS_HEADER) }
    }

    fn nt_headers(&self) -> &IMAGE_NT_HEADERS {
        unsafe {
            &*(self
                .image()
                .as_ptr()
                .offset(self.dos_header().e_lfanew as isize)
                as *mut IMAGE_NT_HEADERS)
        }
    }

    fn file_header(&self) -> &IMAGE_FILE_HEADER {
        &self.nt_headers().FileHeader
    }

    fn opt_header(&self) -> &IMAGE_OPTIONAL_HEADER {
        &self.nt_headers().OptionalHeader
    }

    fn section_headers(&self) -> &[IMAGE_SECTION_HEADER] {
        unsafe {
            let len = self.file_header().NumberOfSections as usize;
            let data = (self.file_header() as *const _ as *const u8)
                .offset(20) // Pe 头固定 20 字节
                .offset(self.file_header().SizeOfOptionalHeader as isize)
                as *const IMAGE_SECTION_HEADER;
            std::slice::from_raw_parts(data, len)
        }
    }

    fn num_of_sections(&self) -> usize {
        self.file_header().NumberOfSections as usize
    }

    fn rva_to_foa(&self, rva: Rva) -> error::Result<u32> {
        // 如果在头部, 直接返回
        if rva < self.opt_header().SizeOfHeaders {
            return Ok(rva);
        }

        // 如果在节中
        let sections = self.section_headers();
        for sec in sections {
            // overflow 情况
            let sec: &IMAGE_SECTION_HEADER = sec;

            let rva_end_bound = sec
                .VirtualAddress
                .wrapping_add(std::cmp::max(sec.SizeOfRawData, sec.VirtualSize));

            // $1
            if rva >= sec.VirtualAddress && rva < rva_end_bound {
                // $2
                if let None = sec.PointerToRawData.checked_add(sec.SizeOfRawData) {
                    return Err(error::Error::Overflow);
                }
                // 不会 underflow, 因为 $1
                let sec_offset = rva - sec.VirtualAddress;
                // $3
                return if sec_offset < sec.SizeOfRawData {
                    // 不会 overflow 因为 $2 $3
                    Ok(sec_offset + sec.PointerToRawData)
                } else if sec_offset < sec.VirtualSize {
                    Err(error::Error::ZeroFill)
                } else {
                    // 不会到这里
                    Err(error::Error::Bounds)
                };
            }
        }
        Err(error::Error::Bounds)
    }

    fn foa_to_rva(&self, foa: u32) -> error::Result<Rva> {
        // headers 中, 直接返回
        if foa < self.opt_header().SizeOfHeaders {
            return Ok(foa);
        }
        // sections 中
        for sec in self.section_headers() {
            let sec: &IMAGE_SECTION_HEADER = sec;
            // 这里不需要比较 SizeOfRawData 和 VirtualSize 谁大
            let foa_end_bound = sec.PointerToRawData.wrapping_add(sec.SizeOfRawData);
            // $1
            if foa >= sec.PointerToRawData && foa < foa_end_bound {
                // overflow $2
                if let None = sec.VirtualAddress.checked_add(sec.VirtualSize) {
                    return Err(error::Error::Overflow);
                }

                // 不会 underflow, 因为 $1
                let sec_offset = foa - sec.PointerToRawData;
                // $3
                return if sec_offset < sec.VirtualSize {
                    // 不会 overflow, 因为 $2 和 $3
                    Ok(sec_offset + sec.VirtualAddress)
                } else if sec_offset < sec.SizeOfRawData {
                    // Unmapped
                    Err(error::Error::Unmapped)
                } else {
                    // 不会到这里
                    Err(error::Error::Bounds)
                };
            }
        }
        Err(error::Error::Bounds)
    }
}
